
// "Copyright [2021] <Copyright Shizhen Zhao>"

#ifndef __INET_MMUMANAGER_H
#define __INET_MMUMANAGER_H

#include <vector>

#include "inet/common/INETDefs.h"

namespace inet {

class INET_API MmuManager {
  public:
	explicit MmuManager(long int num_ports, bool is_infiniMemory_);

    // Configures switch MMU.
    void ConfigMmu(const std::vector<long int>& bandwidth_gbps,
                   const std::vector<long int>& link_delay_ns);

    // Admission control functions for MMU.
	bool CheckIngressAdmission(long int in_port, long int packet_size_byte);
	void UpdateIngressAdmission(long int in_port, long int packet_size_byte);
	void UpdateEgressAdmission(long int out_port, long int packet_size_byte);
	void RemoveFromIngressAdmission(long int in_port, long int packet_size_byte);
	void RemoveFromEgressAdmission(long int out_port, long int packet_size_byte);

	// Checks if ECN should be set based on output queue size.
	bool ShouldSendECN(long int out_port);

    // Uses the following functions to pause or resume. 
	bool CheckShouldPause(long int in_port);
	bool CheckShouldResume(long int in_port);
	void SetPause(long int in_port);
	void SetResume(long int in_port);


	void dispMem();

  private:
    // Helper functions that computes pfc thresholds and shared memory usage.
	bool is_infiniMemory;
	long int GetPfcThreshold();
	long int GetSharedUsed(long int in_port);

	// Switch configurations.
	long int buffer_size_;  // Unit is byte, size of the entire mmu.
    long int resume_offset_;  // Parameter controlling when to resume.
	long int pfc_a_shift_;  // Used to compute pfc threshold.
    long int port_count_;  // Total port count. Each port is bidirectional.

	long int reserve_;  // Each inport reserves a certain amount of memory.
    long int total_rsrv_;  // Total amount of reserved memory.
    // Packets are stored in a headroom vector after triggering pfc.
	std::vector<long int> headroom_;
	long int total_hdrm_;  // Total amount of headroom memory.

	// Variables related to ECN.
	double pmax_;
	std::vector<long int> kmax_;
	std::vector<long int> kmin_;

	// Runtime variables.
	long int shared_used_bytes_;  // Packets exceeding reserve_ memory use shared memory.
	std::vector<long int> hdrm_bytes_;  // Total bytes in the headroom buffer per port
	std::vector<long int> ingress_bytes_;  // Total bytes in each ingress port
	std::vector<long int> egress_bytes_;  // Total bytes in each egress port
	std::vector<bool> paused_;  // If pause signal has been sent from an ingress port
};

}  // namespace inet

#endif // ifndef __INET_MMUMANAGER_H
